#!/usr/bin/python
# -*- coding: utf-8 -*-
#
import os
import re
import sys
#allinone avoid print buffer, instead of use sys.stdout.flush() after each print
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', 0)
#not work: os.environ["PYTHONUNBUFFERED"] = "1"
#
#BSUB -J test_Si
#BSUB -P ase
#BSUB -q testv3
#BSUB -n 4
#BSUB -oo out.%J.log
##BSUB -o jobname.%J.out
##BSUB -e jobname.%J.err
## readin the original hostsfiles of LSF and get hostdict hostdict_str ncpus_tot"
hostfile_org = os.getenv('LSB_DJOB_HOSTFILE')
if (hostfile_org):
    print "\n    original hostsfiles $LSB_DJOB_HOSTFILE is: \n%s"%hostfile_org
    nodelist=[]
    with open(hostfile_org,'r') as f:
        for line in f:
            reline = re.split('[\s]+|[\n]+', line.strip() )
            if reline[0]:
                #print "ok", reline
                nodelist.append(reline[0]) #map(str,line.split('\n')))
            #else:
                #print "not ok", reline
    hostdict = {}
    for host in set(nodelist):
        hostdict[host] = nodelist.count(host)
else:
    hostdict = {'nodeA':8, 'nodeB':8 }
#
hostdict_str = ""
for host in hostdict.keys(): 
    hostdict_str += host+" slots="+str(hostdict[host])+" \n"
ncpus_tot = sum(hostdict.values() )
print "\n    ncpus_tot = %s,  hostdict = %s"%(ncpus_tot, hostdict) 
## save hostlist to file ::
#hostfpath  = "/home/nic/wszhang/eclipse_project/delta_dft/cifs2deltaDFT/myhosts-node45" 
hostfpath  = "myhosts"
host_f = open( hostfpath ,'w') 
host_f.write(hostdict_str) 
host_f.close() 
print "\n    hostfpath = ", hostfpath
print   "    hostlist was saved to $hostfpath as: \n%s \n"%(hostdict_str)
#os.system( "cat $LSB_DJOB_HOSTFILE > "+ hostfpath )


import time
import subprocess
import string
import copy
import numpy as np
import math as math
from io import StringIO
from string import ljust
from _elementtree import Element
#from distutils.cygwinccompiler import get_versions
#from scipy.io.array_import import default

#not work: os.system(" module load ase/3.13.0 ")
#not work: os.system(" module load cp2k/4.1 ") 
#pname = "5layers_fix1layerGa_CO"
import ase
import ase.io as aio
import ase.build as abd
import ase.optimize as aopt
import ase.visualize as av
import ase.constraints as ac
from ase import Atoms
from ase.units import Ry




a0 = 5.43
bulk = Atoms('Si2', [(0, 0, 0),
                     (0.25, 0.25, 0.25)],
             pbc=True)
b = a0 / 2
bulk.set_cell([(0, b, b),
               (b, 0, b),
               (b, b, 0)], scale_atoms=True)


#atoms = aio.read('zngao.cif',format='cif')
#print atoms
##P = np.array([[  3,   0,   0],
##       [  0,   2,   0],
##       [  0,   0,   2]])
##atomss = abd.make_supercell(atoms, P )
#layer  = abd.cut(atoms, (1,0,0), (0,1,-1), nlayers=5, origo=(0,0,-0.02), tolerance=0.015, extend=1.0)
#layers = abd.cut(atoms, (3,0,0), (0,2,-2), nlayers=5, origo=(0,0,-0.02), tolerance=0.015, extend=1.0)
##layer322 = abd.make_supercell(layer, P )
#layers.rotate( (0,1,1), (0,0,1), rotate_cell=True)
#layers.cell[2][2] = 23.
#print layers 
##av.view(layers)


cif_dir = "./cif_crystals"
cif_files = os.listdir( cif_dir )
cif_files = [ fname for fname in cif_files if os.path.isfile(cif_dir+"/"+fname) and ".cif" in fname ]
print cif_files

basis_dir = "./Orb_SG15_DZP_E100_Standard_v1.0"
basis_files = os.listdir( basis_dir )
basis_files = [ fname for fname in basis_files if os.path.isfile(basis_dir+"/"+fname) ]
print basis_files


for cif_name in cif_files: 
    print cif_name
    
    bulk = aio.read('./cif_crystals/'+cif_name, format='cif')

    atom_all = bulk.get_chemical_symbols()
    #print atom_all
    atom_species = []
    [ atom_species.append(ii) for ii in atom_all if ii not in atom_species ]
    #print atom_species

    basis_list = []     
    for i_species in atom_species:
        [ basis_list.append(i_basis) for i_basis in basis_files if i_species+"_" in i_basis ] 
    
    print basis_list
    print " bulk.get_number_of_atoms = %s"%bulk.get_number_of_atoms() 
    print " bulk.get_volume() = %s A^3"%bulk.get_volume()


quit()


from ase.calculators.abacus.abacus_out import *
from ase.calculators.abacus.create_input import *
#from ase.calculators.abacus import abacus_out 
#from abacus_out import Abacus
#from create_input import AbacusInput

os.environ['ASE_ABACUS_COMMAND'] = "mpirun -hostfile ../myhosts /home/nic/wszhang/eclipse_project/abacus_dft/ABACUS.1.0.0/bin/70831_ABACUS.mpi.1.0.0  > PREFIX.log"

#default_parameters = dict(calculation = 'scf',
#                              ecutwfc = 50,
#                              smearing = 'gaussian',
#                              mixing_type = 'pulay-kerker',
#                              basis_type = 'lcao',
#                              gamma_only = 1,
#                              atom_file = 'STRU',
#                              )
#print " Abacus.default_parameters: \n%s\n"%Abacus.default_parameters

calc = Abacus(  label="rundir/test",
                #command = "mpirun -hostfile myhosts /home/nic/wszhang/eclipse_project/abacus-NewGit/ABACUS.1.0.0/bin/20160621ABACUS.mpi.1.0.0 " ,
                #command = "mpirun -hostfile ../myhosts /home/nic/wszhang/eclipse_project/abacus_dft/ABACUS.1.0.0/bin/70831_ABACUS.mpi.1.0.0 " ,
                atoms = bulk,
                potential_path = "" ,
                potential_name = "potential_pbe_sg15" , 
                #basis_path = "/home/nic/wszhang/eclipse_project/delta_results/0abacus_lcao07_E100/orbfoldall_rerun_new/DZP_E100_v1.0/plot_SaveOrb/Orb_SG15_DZP_E100_Standard_v1.0" ,
                basis_path = "../" ,
                #basis_name = [ "O_gga_7au_100Ry_2s2p1d.orb", "Zn_gga_9au_100Ry_4s2p2d1f.orb" ] ,  #"basis_pbe_sg15 accurate" , 
                basis_name = "basis_pbe_sg15 accurate" , 
                pseudo_dir = "/home/nic/wszhang/eclipse_project/delta_dft/CIF_POT/SG15_ONCV_PBE-1.0/",
                ) #**Abacus.default_parameters )



bulk.set_calculator(calc)
#calc.set( default_parameters ) 
#calc.run
#print bulk.get_calculator()
#calc.calculate( atoms=bulk )


e = bulk.get_potential_energy()
print "\n bulk.get_potential_energy() = %s eV (ABACUS)"%e

print "\n bulk.get_forces() \n", bulk.get_forces()

print "\n calc.get_fermi_level() :", calc.get_fermi_level()

quit()
#print bulk.get_kinetic_energy
#print bulk.get_total_energy()
#print calc.read_fermi()
#print calc.results



#from ase.calculators.abinit import Abinit

##os.environ['ABINIT_PP_PATH']=os.path.dirname(os.path.realpath(__file__))
##os.environ['ABINIT_PP_PATH']=os.getcwd()
##os.environ['ABINIT_PP_PATH']="/opt/abinit/JTH-PBE-atomicdata-1.0/"]
#os.environ['ASE_ABINIT_COMMAND'] = "mpirun -hostfile myhosts abinit < PREFIX.files > PREFIX.log"
#print os.environ['ABINIT_PP_PATH']
#print os.environ['ASE_ABINIT_COMMAND']
#
#calc = Abinit(label='Si_test',
#              nbands=60,  # one can specify any abinit keywords
#              ecut=30 * Ry,  # warning - used to speedup the test
#              pawecutdg=120 * Ry,
#              kpts=[6, 6, 6],  # warning - used to speedup the test
#              chksymbreak=0,
#              #autoparal=1,
#              xc='PBE',
#              smearing=['gaussian',0.002 * Ry],
#              pps='pawxml',
#              #command = "   source /opt/abinit/set_env_openmpi200-intel2015.sh" +
#              #          "&& mpirun -np 4 abinit < PREFIX.files > PREFIX.log" ,
#              )
#
#
## one can specify abinit keywords also using set
#calc.set(toldfe=1.0e-5)  # warning - used to speedup the test
#bulk.set_calculator(calc)
#
#e = bulk.get_potential_energy()
#print e
#print bulk.get_total_energy()
#print calc.read_fermi()
#print calc.results

#os._exit(0)
sys.exit(0)



# CP2K
#from ase.calculators.cp2k import CP2K
#CP2K.command = "mpijob cp2k_shell.popt"
